"""Simple MNIST image classifier example with LightningModule and LightningDataModule.

To run: python minmst_example.py 
"""

import torch
import torch_npu

from examples.dataset_moudle.mnist_datamodule import MNISTDataModule
from examples.module.mnist_module import ImageClassifier
from pytorch_lightning.utilities.cli import LightningCLI

from lightning_npu.accelerators.npu import NPUAccelerator
from lightning_npu.strategies.npu_parallel import NPUParallelStrategy


def cli_main():
    # The LightningCLI removes all the boilerplate associated with arguments parsing. This is purely optional.
    cli = LightningCLI(
        ImageClassifier,
        MNISTDataModule,
        trainer_defaults={
            "accelerator": NPUAccelerator(),
            "devices": 8,
            "max_epochs": 5,
            "strategy" : NPUParallelStrategy(),
        },
        seed_everything_default=42,
        save_config_overwrite=True,
        run=False
    )
    cli.trainer.fit(cli.model, datamodule=cli.datamodule)
    cli.trainer.test(ckpt_path="best", datamodule=cli.datamodule)


if __name__ == "__main__":
    cli_main()